# -*- coding: utf-8 -*-
"""
Created on Tue Jun  5 20:34:06 2018

@author: zy
"""

'''
调用Object Detection API进行实物检测   需要GPU运行环境，CPU下会报错

TensorFlow  生成的  .ckpt 和  .pb 都有什么用？
https://www.cnblogs.com/nowornever-L/p/6991295.html
如何用Tensorflow训练模型成pb文件（一）——基于原始图片的读取
https://blog.csdn.net/u011463646/article/details/77918980?fps=1&locationNum=7
'''

# 运行前需要把object_detection添加到环境变量
# ubuntu 在research目录下，打开终端，执行export PYTHONPATH=$PYTHONPATH:${PWD}:${PWD}/slim 然后执行spyder，运行程序
# windows 在research目录下，打开cmd，执行set PYTHONPATH=%PYTHONPATH%;%CD%;%CD%/slim 然后执行spyder，运行程序
import time
import matplotlib.pyplot as plt
import numpy as np
import os
import tensorflow as tf
from object_detection.utils import label_map_util
from object_detection.utils import visualization_utils as vis_util
from PIL import Image


def main():
    # 重置图
    tf.reset_default_graph()
    '''
    载入模型以及数据集样本标签，加载待测试的图片文件
    '''
    # 指定要使用的模型的路径  包含图结构，以及参数
    # PATH_TO_CKPT = r'F:\Resources\model\Detector\frozen_inference_graph.pb'
    PATH_TO_CKPT = r'F:\Resources\model\nouse\object_detection_ssd1000.pb'
    # 测试图片所在的路径
    PATH_TO_TEST_IMAGES_DIR = r'F:\bigphoto\test_images'

    TEST_IMAGE_PATHS = [os.path.join(PATH_TO_TEST_IMAGES_DIR, i) for i in os.listdir(PATH_TO_TEST_IMAGES_DIR)]

    # 数据集对应的label pascal_label_map.pbtxt文件保存了index和类别名之间的映射
    # PATH_TO_LABELS = r'F:\Resources\model\Detector\foxconn_box_label_map.pbtxt'
    PATH_TO_LABELS = r'F:\Resources\model\dete3\foxconn_double_label_map.pbtxt'

    NUM_CLASSES = 1

    # 重新定义一个图
    output_graph_def = tf.GraphDef()
    # print(time.time())

    with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
        # 将*.pb文件读入serialized_graph
        serialized_graph = fid.read()
        # 将serialized_graph的内容恢复到图中
        output_graph_def.ParseFromString(serialized_graph)
        # print(output_graph_def)
        # 将output_graph_def导入当前默认图中(加载模型)
        tf.import_graph_def(output_graph_def, name='')

    print('模型加载完成')

    # 载入coco数据集标签文件
    label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
    categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES,
                                                                use_display_name=True)
    category_index = label_map_util.create_category_index(categories)
    print(category_index)
    '''
    定义session
    '''

    def load_image_into_numpy_array(image):
        '''
        将图片转换为ndarray数组的形式
        '''
        im_width, im_height = image.size
        return np.array(image.getdata()).reshape((im_height, im_width, 3)).astype(np.uint0)

    # 设置输出图片的大小
    IMAGE_SIZE = (12, 8)

    # 使用默认图，此时已经加载了模型
    detection_graph = tf.get_default_graph()

    with tf.Session(graph=detection_graph) as sess:
        for image_path in TEST_IMAGE_PATHS:
            image = Image.open(image_path)
            # 将图片转换为numpy格式
            image_np = load_image_into_numpy_array(image)

            '''
            定义节点，运行并可视化
            '''
            # 将图片扩展一维，最后进入神经网络的图片格式应该是[1,?,?,3]
            image_np_expanded = np.expand_dims(image_np, axis=0)

            '''
            获取模型中的tensor
            '''
            image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')

            # boxes用来显示识别结果
            boxes = detection_graph.get_tensor_by_name('detection_boxes:0')

            # Echo score代表识别出的物体与标签匹配的相似程度，在类型标签后面
            scores = detection_graph.get_tensor_by_name('detection_scores:0')
            classes = detection_graph.get_tensor_by_name('detection_classes:0')
            num_detections = detection_graph.get_tensor_by_name('num_detections:0')

            # 开始检查
            print(time.strftime('%Y/%m/%d %H:%M:%S', time.localtime(time.time())))
            # print(time.time())

            boxes, scores, classes, num_detections = sess.run([boxes, scores, classes, num_detections],
                                                              feed_dict={image_tensor: image_np_expanded})
            print('score',scores)
            print(time.strftime('%Y/%m/%d %H:%M:%S', time.localtime(time.time())))
            # print(time.time())
            # 可视化结果
            # print('classs',classes)
            vis_util.visualize_boxes_and_labels_on_image_array(
                image_np,
                np.squeeze(boxes),
                np.squeeze(classes).astype(np.int32),
                np.squeeze(scores),
                category_index,
                use_normalized_coordinates=True,
                line_thickness=4)
            plt.figure(figsize=IMAGE_SIZE)
            print(type(image_np))
            print(image_np.shape)
            image_np = np.array(image_np, dtype=np.uint8)
            plt.imshow(image_np)
            plt.show()

if __name__ == '__main__':
    main()